fix(kernels): drop the D2H sync from the varlen GDN/KDA prefill conv - #339

Thanks for this one — the sync you are removing is real, and the reasoning in the description
matches what I see in the tree. I pulled the branch onto an Ada card to try it and the fix itself
is fine; what I ran into is that the three new tests all fail here, and the cause is that they
assume the triton fallback is the active path. Since sgl is an optional extra, I suspect the
tests pass on your install and fail on any freetoken[accel] one, which would make this hard to
land without someone hitting it.

Setup. RTX 4090 24 GB (sm_89), CUDA toolkit 13.3, sglang-kernel 0.4.5 installed
(freetoken[accel]). This PR on top of main @ 03c28d2 plus PR #337 @ a6bd5c0. Every number
below is from this box.

tests/kernels/test_causal_conv1d_capture.py    3 failed, 2 passed

Cause 1 — the two backends disagree about whether x is mutated (2 of the 3)

causal_conv1d_varlen dispatches on is_sgl_kernel_installed(), and the two branches differ:

Measured by calling _call() from your own _conv_inputs() with the dispatch forced each way.

test_varlen_conv_with_host_metadata_matches_the_device_derived_result and
test_varlen_conv_replays_inside_a_cuda_graph both restore conv_states between the two calls
they compare, but not x. On the triton path that is correct, because x is untouched. With
sgl_kernel the second call convolves an already-convolved x, so the two results differ and the
torch.equal assertions fail. Nothing is wrong with the kernel or with your change here — the
comparison is just not starting from the same input twice.

Cause 2 — the D2H sync only exists on the fallback path (the third)

test_varlen_conv_still_derives_max_seq_len_on_device_by_default asserts that .item() is called
when max_seq_len is omitted. Counting torch.Tensor.item calls with the dispatch forced each
way:

sgl_kernel path     .item() called 0 times   -> assertion fails
triton fallback     .item() called 1 time    -> assertion holds

Which is exactly right: int(seq_lens.max().item()) lives in
kernel/triton/causal_conv1d_triton.py, and the native kernel does not need the value at all. So
the test is asserting a property of the fallback while running whichever backend happens to be
installed.

Suggested test fix — verified, 5 passed here

Restore x alongside conv_states, and pin the third test to the path whose behaviour it
describes:

     baseline_states = inputs["conv_states"].clone()
+    baseline_x = inputs["x"].clone()
...
     inputs["conv_states"].copy_(baseline_states)
+    inputs["x"].copy_(baseline_x)

+    # The sync this asserts on lives in the triton fallback; on an install with
+    # sgl_kernel the native kernel needs no max_seq_len and calls no .item().
+    import freetoken.kernel.backend as _backend
+    monkeypatch.setattr(_backend, "is_sgl_kernel_installed", lambda: False)
     monkeypatch.setattr(torch.Tensor, "item", counted_item)

With both applied: 5 passed on this box. The second one is the part I would argue for on its
own merits — forcing the fallback means the test proves the claim on any install, rather than
only where the fallback happens to be selected. (A skipif(is_sgl_kernel_installed()) would also
go green, but it would stop testing the thing on the machines most likely to run CI.)

I have not opened this as a PR; the patch is small enough to paste, and it is your branch. Happy
to send it if you would rather have it that way.

One separate observation, offered as a note rather than a request

The x-mutation divergence above is not caused by this PR and is harmless in the tree today: all
three call sites (qwen3_5_moe/gdn.py:122, qwen4_exp/gdn.py:131, glm5_next/kda.py:183) build
x = conv_in.transpose(0, 1).contiguous() immediately before the call and never read it again, so
nobody depends on x surviving. But the wrapper's docstring does not say which contract holds,
and a future caller that keeps x would break on one install shape and not the other — the same
way these tests just did. Might be worth a line in the wrapper's docstring while this file is
open. I have not audited beyond the three call sites above.

For what it is worth, the fix is a no-op on my own serving path for the same reason — with
sgl_kernel installed the sync never executes — so I cannot give you a before/after timing. On a
default install (no [accel]), where the fallback is the path, the change should do exactly
what you describe.

_{Written with AI assistance; every number above was measured on my hardware (RTX 4090, sm_89)
and I can reproduce it.}

Good catch, and the diagnosis matches what I see in the tree: the sgl path returns x mutated in place, the fallback returns a fresh tensor, and the .item() only lives in causal_conv1d_triton.py. I only have the fallback here (no sgl_kernel on this Windows box), which is why the tests looked fine on my side.

Pushed eea4ad6 with your two changes plus the docstring note:

• x is restored alongside conv_states between the compared calls (both tests)
• the "derives max_seq_len on device" test now forces is_sgl_kernel_installed to False, for the reason you gave — it should prove the claim on any install, not just where the fallback happens to be selected
• the wrapper docstring now says which contract holds on which backend, and that callers must use the returned tensor rather than x

To cover the path I can't run natively, I also ran the file with a stand-in sgl_kernel whose causal_conv1d_fwd convolves x in place and returns it. The unmodified tests fail 3/5 under it (same three you listed), the pushed version passes 5/5, and it still passes 5/5 on the real fallback. If you get a chance to re-run on the 4090 that would be the real confirmation.

Confirmed on the real thing — eea4ad6 passes here.

RTX 4090 (sm_89), sglang-kernel 0.4.5 installed, is_sgl_kernel_installed() -> True
  tests/kernels/test_causal_conv1d_capture.py     5 passed

For the before/after on this box: the previous head failed 3 of 5 here (the three I listed),
and this one is 5/5. So your stand-in sgl_kernel reproduced the real backend's behaviour
exactly — same three tests, same direction. That was a good way to reach a path you cannot run.

Also ran the GDN callers that go through the wrapper, to be sure nothing shifted underneath:
tests/models/qwen4_exp/test_gdn.py 6 passed.

The docstring addition is the part I would have argued for hardest, so I am glad you took it:

Whether x survives depends on the backend: sgl_kernel convolves x in place and returns it,
while the triton fallback leaves x untouched and returns a new tensor. Callers must use the
returned tensor and must not rely on x afterwards

That turns the thing that broke the tests into a stated contract, which is the part that will
outlive this PR once it lands. The three current call sites all build x fresh and never read it again, so nothing has
to change today — but the next one now has something to read.

Nothing further from me on this one. Happy to re-run on the 4090 if the branch moves again.

_{Written with AI assistance; the test results above are from my own hardware and I can
reproduce them.}

Tested on a Qwen3.8-Flash-Next deployment; no measurable change, and no regression.

Caveat on what this measured: sgl_kernel is installed on this box, so causal_conv1d_varlen took the sgl_kernel path and the .item() this PR removes from the triton fallback never ran here. The numbers only say the host-side max_seq_len plumbing through FLAMetadata costs nothing on the fused path. On this model the 1k-token TTFT is dominated by expert streaming over PCIe, so even on the fallback path one sync per GDN layer per prefill would be hard to see in TTFT; the capture-ability argument is the real value.

Measured on a 2x RTX 6000 Ada (48 GiB, sm_89, PCIe Gen4, no NVLink) / 2x Xeon Gold 6526Y / 503 GiB box, CUDA 13.3, torch 2.11+cu130, sgl_kernel 0.4.5, model RadixArk/Qwen3.8-Flash-Next-NVFP4 (qwen4_exp, NVFP4 experts + bf16 dense). The PR cherry-picks cleanly onto main 86214a9; both runs also carry a local fix so --moe-cache-auto honours --num-tokens (it OOMs otherwise at this KV size). One GPU, one run each, same flags:

ft serve --model <RadixArk snapshot> --moe-backend offload --ple-backend pinned --num-tokens 262144 \
  --memory-ratio 0.94 --moe-prefill-hit-d2d --max-running-requests 8 --cuda-graph-max-bs 8

Single-stream = median of three 64-vs-256-token completion pairs, aggregate = 8 concurrent 256-token completions, TTFT on a ~1k-token prompt. Run-to-run spread of the baseline on this box is about +-4% single-stream.